import mujoco.viewer
import numpy as np

from src.mj_interface import MJInterface
from src.utils.transformation import Transformation
from src.robotics.robotics import Robotics
from src.robot_state import RobotState, MotionState
from src.wbc import WBC
from src.mpc import MPC
from src.gait_scheduler import GaitScheduler, SwingPlanner
from src.controller import PVTController
from src.joystick import Joystick

if __name__ == '__main__':
    mj_model: mujoco.MjModel = mujoco.MjModel.from_xml_path("../assets/humanoid-legged-description/mjcf/scene.xml")
    mj_data = mujoco.MjData(mj_model)

    dt_200Hz = 0.005

    mj_interface = MJInterface(mj_model, mj_data)
    robotics = Robotics('../assets/humanoid-legged-description/mjcf/humanoid_legged.urdf')
    fe_rot = np.array([
        [0, 0, 1],
        [1, 0, 0],
        [0, 1, 0]
    ])
    robot_state = RobotState(robotics.model_nv, fe_rot=fe_rot)
    wbc = WBC(robotics.model_nv, 18, 22, 0.5, mj_model.opt.timestep)
    mpc = MPC(dt_200Hz)
    gait_scheduler = GaitScheduler(robotics.model_nv, 0.2, mj_model.opt.timestep)
    pvt_controller = PVTController(mj_model.opt.timestep, '../src/common/joint_ctrl_config.yaml')
    swing_planner = SwingPlanner()
    joystick = Joystick(mj_model.opt.timestep)

    stand_legLength = 1.19
    foot_height = 0.07
    xv_des = 0.2
    wz_des = 0.0

    robot_state.width_hips = 0.2995
    swing_planner.kp_vx = 0.5
    swing_planner.kp_vy = 0.4
    swing_planner.kp_wz = 0.03
    swing_planner.stepHeight = 0.2
    swing_planner.legLength = stand_legLength

    model_nv = robotics.model_nv

    motors_pos_des = np.zeros(model_nv - 6)
    motors_pos_cur = np.zeros(model_nv - 6)
    motors_vel_des = np.zeros(model_nv - 6)
    motors_vel_cur = np.zeros(model_nv - 6)
    motors_tau_des = np.zeros(model_nv - 6)
    motors_tau_cur = np.zeros(model_nv - 6)
    fe_l_pos_L_des = np.array([0.083, 0.16, -stand_legLength])
    fe_r_pos_L_des = np.array([0.083, -0.16, -stand_legLength])
    fe_l_eul_L_des = np.array([-0.000, -0.008, -0.000])
    fe_r_eul_L_des = np.array([0.000, -0.008, 0.000])
    fe_l_rot_des = Transformation.eul2Rot(fe_l_eul_L_des[0], fe_l_eul_L_des[1], fe_l_eul_L_des[2]) @ robot_state.fe_rot
    fe_r_rot_des = Transformation.eul2Rot(fe_r_eul_L_des[0], fe_r_eul_L_des[1], fe_r_eul_L_des[2]) @ robot_state.fe_rot

    resLeg = robotics.computeInk_Leg(fe_l_rot_des, fe_l_pos_L_des, fe_r_rot_des, fe_r_pos_L_des)

    qIniDes = np.zeros(mj_model.nq)
    qIniDes[7:] = resLeg.jointPosRes
    wbc.setQini(qIniDes, robot_state.q)
    MPC_count = 0

    simEndTime = 50.0
    simTime = mj_data.time
    startSteppingTime = 1.0
    startWalkingTime = 1.2

    with mujoco.viewer.launch_passive(mj_model, mj_data) as viewer:
        while viewer.is_running():
            mujoco.mj_step(mj_model, mj_data)

            simTime = mj_data.time
            mj_interface.updateSensorValues()
            mj_interface.write_state(robot_state)

            robotics.read_state(robot_state)
            robotics.computeJ_dJ()
            robotics.computeDyn()
            robotics.write_state(robot_state)

            if simTime > startWalkingTime:
                joystick.set_wz_des_L_para(wz_des, 1.0)
                joystick.set_vx_des_L_para(xv_des, 2.0)
                robot_state.motionState = MotionState.WALK
            else:
                joystick.setIniPos(robot_state.q[0], robot_state.q[1], robot_state.base_rpy[2])

            joystick.step()
            robot_state.js_pos_des[2] = stand_legLength + foot_height
            joystick.write_state(robot_state)

            if simTime >= startSteppingTime:
                gait_scheduler.read_state(robot_state)
                gait_scheduler.step()
                gait_scheduler.write_state(robot_state)

                swing_planner.read_state(robot_state)
                swing_planner.getSwingPos()
                swing_planner.write_state(robot_state)

            MPC_count += 1
            if MPC_count > (dt_200Hz / mj_model.opt.timestep - 1):
                mpc.read_state(robot_state)
                mpc.calculate()
                mpc.write_state(robot_state)
                MPC_count = 0

            wbc.read_state(robot_state)
            wbc.handle_task(robotics)
            wbc.calculate()
            wbc.write_state(robot_state)

            if simTime <= startSteppingTime:
                robot_state.motors_pos_des[:] = resLeg.jointPosRes
                robot_state.motors_vel_des[:] = motors_vel_des[:]
                robot_state.motors_tor_des[:] = motors_tau_des[:]
            else:
                mpc.enable()
                L_diag = np.array([1.0, 1.0, 1.0,
                                   1.0, 200.0, 1.0,
                                   1e-7, 1e-7, 1e-7,
                                   100.0, 10.0, 1.0])
                K_diag = np.array([1.0, 1.0, 1.0,
                                   1.0, 1.0, 1.0,
                                   1.0, 1.0, 1.0,
                                   1.0, 1.0, 1.0, 1.0])
                mpc.set_weight(1e-6, L_diag, K_diag)

                pos_des = robotics.integrateDIY(robot_state.q, robot_state.wbc_delta_q_final)
                robot_state.motors_pos_des[:] = pos_des[7:].copy()
                robot_state.motors_vel_des[:] = robot_state.wbc_dq_final[6:]
                robot_state.motors_tor_des[:] = robot_state.wbc_tauJointRes[:]

            pvt_controller.read_state(robot_state)
            if simTime <= startSteppingTime:
                pvt_controller.calMotorsPVT(100.0 / 1000.0 / 180.0 * 3.1415)
            else:
                pvt_controller.setJointPD(20.0, 2.0, "leg_l5_joint")
                pvt_controller.setJointPD(20.0, 2.0, "leg_r5_joint")
                pvt_controller.setJointPD(20.0, 2.0, "leg_l6_joint")
                pvt_controller.setJointPD(20.0, 2.0, "leg_r6_joint")
                pvt_controller.setJointPD(500.0, 50.0, "leg_l4_joint")
                pvt_controller.setJointPD(500.0, 50.0, "leg_r4_joint")
                pvt_controller.calMotorsPVT()
            pvt_controller.write_state(robot_state)

            mj_interface.setMotorsTorque(robot_state.motors_tor_out)

            if mj_data.time >= simEndTime:
                break

            viewer.sync()
